// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//


#include <optional>
#include <pollux/testing/gtest_utils.h>
#include <pollux/testing/functions/prestosql/function_base_test.h>

using namespace kumo::pollux;
using namespace kumo::pollux::test;
using namespace kumo::pollux::functions::test;

namespace {

class ArraySumTest : public FunctionBaseTest {
 protected:
  // Evaluate an expression.
  template <typename T>
  void testArraySum(const VectorPtr& expected, const VectorPtr& input) {
    auto result = evaluate("array_sum(c0)", make_row_vector({input}));
    assertEqualVectors(expected, result);

    // Test constant input.
    const vector_size_t firstRow = 0;
    result = evaluate(
        "array_sum(c0)",
        make_row_vector({BaseVector::wrap_in_constant(10, firstRow, input)}));
    assertEqualVectors(
        BaseVector::wrap_in_constant(10, firstRow, expected), result);

    const vector_size_t lastRow = input->size() - 1;
    result = evaluate(
        "array_sum(c0)",
        make_row_vector({BaseVector::wrap_in_constant(10, lastRow, input)}));
    assertEqualVectors(
        BaseVector::wrap_in_constant(10, lastRow, expected), result);
  }
};

} // namespace

// Test integer arrays.
TEST_F(ArraySumTest, int64Input) {
  auto input = make_nullable_array_vector<int64_t>(
      {{0, 1, 2}, {std::nullopt, 1, 2}, {std::nullopt}});
  auto expected = make_nullable_flat_vector<int64_t>({3, 3, 0});
  testArraySum<int64_t>(expected, input);
}

TEST_F(ArraySumTest, int32Input) {
  auto input = make_nullable_array_vector<int32_t>(
      {{0, 1, 2}, {std::nullopt, 1, 2}, {std::nullopt}});
  auto expected = make_nullable_flat_vector<int64_t>({3, 3, 0});
  testArraySum<int64_t>(expected, input);
}

TEST_F(ArraySumTest, int16Input) {
  auto input = make_nullable_array_vector<int16_t>(
      {{0, 1, 2}, {std::nullopt, 1, 2}, {std::nullopt}});
  auto expected = make_nullable_flat_vector<int64_t>({3, 3, 0});
  testArraySum<int64_t>(expected, input);
}

TEST_F(ArraySumTest, int8Input) {
  auto input = make_nullable_array_vector<int8_t>(
      {{0, 1, 2}, {std::nullopt, 1, 2}, {std::nullopt}});
  auto expected = make_nullable_flat_vector<int64_t>({3, 3, 0});
  testArraySum<int64_t>(expected, input);
}

TEST_F(ArraySumTest, overflow) {
  // Flat input.
  auto input = make_row_vector({
      make_nullable_array_vector<int64_t>(
          {{1, 2, 3}, {0, std::numeric_limits<int64_t>::max(), 2, 3}, {}}),
  });
  POLLUX_ASSERT_THROW(
      evaluate("array_sum(c0)", input),
      "integer overflow: 9223372036854775807 + 2");

  auto result = evaluate("try(array_sum(c0))", input);
  assertEqualVectors(
      make_nullable_flat_vector<int64_t>({6, std::nullopt, 0}), result);

  // Constant input.
  auto emptyInput = make_row_vector(ROW({}), 3);
  POLLUX_ASSERT_THROW(
      evaluate(
          "array_sum(array[9223372036854775807, 100::bigint])", emptyInput),
      "integer overflow: 9223372036854775807 + 100");

  result = evaluate(
      "try(array_sum(array[9223372036854775807, 100::bigint]))", emptyInput);
  assertEqualVectors(
      make_nullable_flat_vector<int64_t>(
          {std::nullopt, std::nullopt, std::nullopt}),
      result);
}

// Test floating point arrays
TEST_F(ArraySumTest, realInput) {
  auto input = make_nullable_array_vector<float>(
      {{0, 1, 2}, {std::nullopt, 1, 2}, {std::nullopt}});
  auto expected = make_nullable_flat_vector<double>({3, 3, 0});
  testArraySum<double>(expected, input);
}

TEST_F(ArraySumTest, doubleInput) {
  auto input = make_nullable_array_vector<double>(
      {{0, 1, 2}, {std::nullopt, 1, 2}, {std::nullopt}});
  auto expected = make_nullable_flat_vector<double>({3, 3, 0});
  testArraySum<double>(expected, input);
}

TEST_F(ArraySumTest, doubleInputLimits) {
  auto input = make_nullable_array_vector<double>(
      {{0, std::numeric_limits<double>::infinity(), 2},
       {std::numeric_limits<double>::quiet_NaN(), 1, 2},
       {std::numeric_limits<double>::lowest(), -1},
       {std::numeric_limits<double>::max(), 1.0}});
  auto expected = make_nullable_flat_vector<double>(
      {std::numeric_limits<double>::infinity(),
       std::numeric_limits<double>::quiet_NaN(),
       std::numeric_limits<double>::lowest(),
       std::numeric_limits<double>::max()});
  testArraySum<double>(expected, input);
}
